import os
from dotenv import load_dotenv 
from langchain_deepseek import ChatDeepSeek
from langgraph.prebuilt import create_react_agent
from langchain_core.tools import tool
from pydantic import BaseModel, Field
import json
import pymysql

# 加载环境变量
load_dotenv(override=True)

# ✅ 创建自然语言转Cypher的工具
class ExtractQuerySchema(BaseModel):
    cypher_query: str = Field(description="用于从 图数据库 提取数据的 Cypher 查询语句。")

# 注册为 Agent 工具
@tool(args_schema=ExtractQuerySchema)
def text_to_cypher(cypher_query) -> str:
    """
    用于将用户的语言描述转化为可执行的 Cypher 查询语句。
    :param cypher_query: 字符串形式的SQL查询语句，用于提取MySQL中的某张表。
    :return：格式化的cypher_query表语句
    """
    return cypher_query

# ✅ 创建SQL查询工具
description = """
当用户需要进行数据库查询工作时，请调用该函数。
该函数用于在指定MySQL服务器上运行一段SQL代码，完成数据查询相关工作，
并且当前函数是使用pymsql连接MySQL数据库。
本函数只负责运行SQL代码并进行数据查询，若要进行数据提取，则使用另一个extract_data函数。
"""

# 定义结构化参数模型
class SQLQuerySchema(BaseModel):
    sql_query: str = Field(description=description)

# ✅ 创建提示词模板
prompt = """
你是一名经验丰富的智能数据分析助手，擅长帮助用户高效完成以下任务：
将自然语言转换为可执行的Cypher：
   图数据库的Schema如下：
    节点类型:
    1. 自然人(Person) 
       - 属性: 
         - id (String): 唯一标识
         - name (String): 姓名
         - age (Integer): 年龄
    
    2. 手机号(Phone)
       - 属性:
         - number (String): 手机号码
         - carrier (String): 运营商
    
    3. 套餐(Plan)
       - 属性:
         - name (String): 套餐名称
         - monthlyFee (Float): 月费
         - dataLimit (Integer): 流量限制(GB)
    
    关系类型:
    1. 拥有(OWNS) 
       - 方向: (Person)-[OWNS]->(Phone)
       - 属性:
         - since (Date): 拥有起始日期
    
    2. 办理(SUBSCRIBED_TO)
       - 方向: (Phone)-[SUBSCRIBED_TO]->(Plan)
       - 属性:
         - startDate (Date): 套餐开始日期
         - endDate (Date): 套餐结束日期

**回答要求：**
- 所有回答均使用**简体中文**，清晰、礼貌、简洁。
- 如果调用工具返回结构化JSON数据，你应提取其中的关键信息简要说明，并展示主要结果。
- 若需要用户提供更多信息，请主动提出明确的问题。

**风格：**
- 专业、简洁、以数据驱动。
- 不要编造不存在的工具或数据。

请根据以上原则为用户提供精准、高效的协助。
"""

# ✅ 创建工具列表
tools = [text_to_cypher]

# ✅ 创建模型
model = ChatDeepSeek(model="deepseek-chat")

# ✅ 创建图 （Agent）
graph = create_react_agent(model=model, tools=tools, prompt=prompt)